import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
import cupy as cp
import pprint

# SVD
d4 = cp.linspace(1, 1000, 1000).reshape(10, 4, 5, 5)  # [N,C,d,d]

# N, C, d, _ = (10, 4, 5, 5)
# d2 = cp.empty([d * N, d * C])


def four_2_two(d4):
    N, C, d, _ = d4.shape
    d2 = cp.empty([d * N, d * C])
    def innrt(i1, i2, i3, i4, d):
        j1 = (i1 - 1) * d + i3
        j2 = (i2 - 1) * d + i4
        return j1, j2
    for i1 in range(N):
        for i2 in range(C):
            for i3 in range(d):
                for i4 in range(d):
                    j1, j2 = innrt(i1, i2, i3, i4, d)
                    d2[j1, j2] = d4[i1, i2, i3, i4]
    return d2
re_d2 =four_2_two(d4)
print(len(cp.unique(re_d2)))


def two_2_four(d2):
    h,w = d2.shape
    N, C, d, _ = (10, 4, 5, 5)
    d4 = cp.empty((10, 4, 5, 5))
    def innrt(i1, i2, i3, i4, d):
        j1 = (i1 - 1) * d + i3
        j2 = (i2 - 1) * d + i4
        return j1, j2
    for i1 in range(N):
        for i2 in range(C):
            for i3 in range(d):
                for i4 in range(d):
                    j1, j2 = innrt(i1, i2, i3, i4, d)
                    # print(i1,i2,i3,i4)
                    # print(j1,j2)
                    d4[i1, i2, i3, i4] = d2[j1, j2]

    return d4

re_d4= two_2_four(re_d2)
print(len(cp.unique(re_d4)))
# U,Sigma, VT = cp.linalg.svd(d2)
# print('U: ',U)
# print('Sigma : ',Sigma)
# print('VT: ',VT)
